import os
import cv2
import time
import argparse
import torch
import warnings
import numpy as np
import math
from yolo import YOLO
import warnings
from deep_sort import build_tracker
# from utils.draw import draw_boxes
from utils.parser import get_config
from utils.draw import draw_boxes,pil_draw_box,pil_draw_box_2
from trafficLine import *
from utils.utils import *
from licence import Licence
import shutil
from detectCar.test import detectCar
from PyQt5.QtGui import *
import json
# from utils.log import get_logger
# from utils.io import write_results

# detect.py
import cv2
import numpy as np
from PIL import ImageFont

from detect2 import VideoTracker as BaseVideoTracker
from utils.draw import pil_draw_box
from utils.parser import get_config
from utils.utils import plot_one_box
from yolo import YOLO
from shared import shared_data, shared_data_lock
import warnings
from datetime import datetime
#后端链接
import requests

class VideoTracker(BaseVideoTracker):
    def __init__(self, cfg, yoloweight="./weights/yolov5s.pt"):
        super().__init__(cfg)
        self._custom_output_dir_set = True  # 添加自定义目录已设置标志已只读

        self.cfg = cfg
        self.detector = YOLO(yoloweight)  # 直接使用参数传入的路径
        # 设置 outputDir 的默认值为前端根目录下的 output/camera_0
        self.outputDir = r"D:\GJ\Azuoye\Axsq\third\A\traffic_work\vue-project1\public\images\violations"
        print(f"[INFO] 输出目录: {self.outputDir}")
        self.uploaded_ids = set()  # 已上传违法信息的ID集合
        # 初始化检测跟踪所需变量
        self.carLocation1 = {}
        self.carLocation2 = {}
        self.carSpeed = {}
        self.carDirection = {}
        self.carPre = {}
        self.virtureLine = [0, 0, 0, 0]
        self.inCar = set()
        self.outCar = set()
        self.trafficLine = []
        self.idx_frame = 0
        self.carLicense = {}
        self.licence = Licence()
        self.detectCar = detectCar()
        self.frameAll = 0
        self.vout = cv2.VideoWriter()
        self.carLabel = {}
        self.saveVideoFlag = True
        self.displayFlag = True
        self.videoFps = 30
        self.videoHeight = 0
        self.trafficLine1 = []
        self.trafficLine2 = []
        self.videoWidth = 0
        self.carSpeedLimit = 0
        self.trafficJamLimit = 0
        self.yolov5weight = ""
        self.carFromLeft = {}
        self.carFromRight = {}
        self.carFromForward = {}
        self.carTurn = {}
        self.endFlag = False
        self.trafficLightFlag = False
        self.dir = ['UP', 'RIGHT', 'DOWN', "LEFT"]
        self.font = ImageFont.truetype(font='model_data/simhei.ttf',
                                       size=np.floor(0.012 * (1920)).astype('int32'))  # 预设大小

        use_cuda = torch.cuda.is_available()
        if not use_cuda:
            warnings.warn("Running in cpu mode which maybe very slow!", UserWarning)

        self.deepsort = build_tracker(cfg, use_cuda=use_cuda)


    def process_frame(self, frame, config):

        """
        处理单帧图像并返回结果图像
        """
        # 初始化
        if not hasattr(self, 'initialized') or not self.initialized:
            self._init_for_stream(config)

        self.idx_frame += 1

        im1 = frame.copy()

        ori_im = frame.copy()  # 原始帧，用于保存违法截图时避免绘制线框
        draw_im = frame.copy()  # 绘制帧，用于返回给前端显示


        # 设置 trafficLine，默认使用 config.trafficline1
        # 使用 config.trafficline1，若不存在或不合法则使用默认值或跳过处理
        if hasattr(config, 'trafficline1') and len(config.trafficline1) >= 4:
            traffic_line = config.trafficline1
        else:
            traffic_line = [0,0,999,999]
        (bbox_xywh, cls_conf, cls_ids, draw_im) = self.detector.detect_image(im1, traffic_line, "", 0, {}, config)
        outputs = self.deepsort.update(bbox_xywh, cls_conf, im1)

        # 如果有输出结果，进行车辆追踪与违法判断
        if len(outputs) > 0:
            bbox_xyxy = outputs[:, :4]
            identities = outputs[:, -1]

            for i, bbox in enumerate(bbox_xyxy):
                if len(bbox) < 4:
                    print(f"[WARN] BBox length invalid: {bbox}")
                    continue

                x1, y1, x2, y2 = bbox_xyxy[i]
                id = int(identities[i]) if identities is not None else 0



                # 更新位置信息
                if self.carPre.get(id, 0) == 0:
                    self.carPre[id] = self.idx_frame
                    self.carLocation1[id] = [x1, y1, x2, y2]
                elif (self.idx_frame - self.carPre[id]) >= 12:
                    self.carLocation2[id] = [x1, y1, x2, y2]

                    loc1 = self.carLocation1.get(id, [])
                    loc2 = self.carLocation2.get(id, [])

                    if not loc1 or not loc2 or len(loc1) < 4 or len(loc2) < 4:
                        continue

                    # 计算速度方向
                    self.carSpeed[id] = self.calculateSpeed(
                        self.carLocation1[id], self.carLocation2[id],
                        self.idx_frame - self.carPre[id]
                    )

                    if self.carSpeed[id] >= 10:
                        self.carDirection[id] = self.dir[self.calculateDirection(
                            self.carLocation1[id], self.carLocation2[id]
                        )]

                    # 判断超速
                    if self.carSpeed.get(id, 0) > self.carSpeedLimit:
                        self.save_illegal_speed(id, x1, y1, x2, y2, frame)

                    # 判断未按导向行驶
                    pre_dir = self.carDirection.get(id, "")
                    if pre_dir == 'UP' and self.carDirection.get(id, "") in ['LEFT', 'RIGHT']:
                        self.carTurn[id] = self.carDirection[id]
                    if self.carTurn.get(id, "") == 'LEFT' and \
                       (self.carFromRight.get(id, False) or self.carFromForward.get(id, False)) and \
                       not self.carFromLeft.get(id, False):
                        self.save_illegal_turnwrong(id, x1, y1, x2, y2, frame)
                    elif self.carTurn.get(id, "") == 'RIGHT' and \
                         (self.carFromLeft.get(id, False) or self.carFromForward.get(id, False)) and \
                         not self.carFromRight.get(id, False):
                        self.save_illegal_turnwrong(id, x1, y1, x2, y2, draw_im)

                    # 判断压线
                    if self.carDirection.get(id, "") == 'UP' and len(config.laneline) > 0:
                        if self.judge_line_illegal(id, [x1, y1, x2, y2], config):
                            self.save_illegal_lane(id, x1, y1, x2, y2, frame)

                    # 判断闯红灯
                    self.carrunred(id, x1, x2, y1, y2, frame)

                    # 更新坐标与时间戳
                    self.carLocation1[id] = self.carLocation2[id].copy()
                    self.carPre[id] = self.idx_frame

                    # 更新车辆信息
                    self.carInfor[id] = f"{self.carDirection.get(id, '')},{self.carLicense.get(id, '')}," \
                                        f"{int(self.carSpeed.get(id, 0))},{self.carLabel.get(id, '')}"

                # 更新车牌与车型
                if self.carLicense.get(id, 0) == 0:
                    t = self.licence.detectLicence(draw_im[y1:y2, x1:x2], x1, y1)
                    if t:
                        xyxy, label = t
                        self.carLicense[id] = label

                if self.carLabel.get(id, 0) == 0 and config.CARCLASSIFY:
                    t = self.detectCar.detect(draw_im[y1:y2, x1:x2])
                    if t:
                        self.carLabel[id] = t

                # 区域判断（左/右/直行）
                w = x2 - x1
                h = y2 - y1
                cx = x1 + w / 2
                cy = y1 + h / 2
                if len(config.leftline) and self.inArea(cx, cy, config.leftline):
                    self.carFromLeft[id] = True
                if len(config.rightline) and self.inArea(cx, cy, config.rightline):
                    self.carFromRight[id] = True
                if len(config.forwardline) and self.inArea(cx, cy, config.forwardline):
                    self.carFromForward[id] = True

            # 绘制车辆信息
            draw_im = pil_draw_box(draw_im, bbox_xyxy, identities,
                                   self.carSpeed, self.carLicense,
                                   self.carLabel, self.carDirection, self.font)
            font = ImageFont.truetype(font='model_data/simhei.ttf',
                                      size=np.floor(0.012 * np.shape(draw_im)[1]).astype('int32'))


            # 显示 carIn / carOut 数量
            draw_im = cv2.putText(draw_im, f"carIn: {len(self.inCar)} carOut: {len(self.outCar)}",
                                  (0, 150), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 255, 0), 2)

            # 绘制车道线
            if len(self.trafficLine1) > 0:
                plot_one_box(self.trafficLine1, draw_im, label="TrafficLine1", color=(205, 210, 104), line_thickness=3)
            if len(self.trafficLine2) > 0:
                plot_one_box(self.trafficLine2, draw_im, label="TrafficLine2", color=(205, 210, 104), line_thickness=3)



        # 共享内存更新
        with shared_data_lock:
            shared_data.frame = draw_im.copy()


        return draw_im

    def _init_for_stream(self, config):
        self.idx_frame = 0
        self.initialized = True
        self.inCar = set()
        self.outCar = set()
        self.init_output_dirs()
        self.carLocation1 = {}
        self.carLocation2 = {}
        self.carSpeed = {}
        self.carDirection = {}
        self.carLicense = {}
        self.illegal = {}
        self.carSpeedLimit = config.CARSPEEDLIMIT
        print("Initialized for streaming")

    def init_output_dirs(self):
        os.makedirs(os.path.join(self.outputDir, "illegal", "overspeed"), exist_ok=True)
        os.makedirs(os.path.join(self.outputDir, "illegal", "turnwrong"), exist_ok=True)
        os.makedirs(os.path.join(self.outputDir, "illegal", "touchline"), exist_ok=True)
        os.makedirs(os.path.join(self.outputDir, "illegal", "carrunred"), exist_ok=True)
        os.makedirs(os.path.join(self.outputDir, "illegal", "runred"), exist_ok=True)

    def calculateSpeed(self, location1, location2, cnt, flag=0):
        x11, y11, x12, y12 = location1
        x21, y21, x22, y22 = location2

        w1 = abs(x12 - x11)
        h1 = abs(y12 - y11)
        w2 = abs(x22 - x21)
        h2 = abs(y22 - y21)
        if w1 <= 0 or h1 <= 0 or w2 <= 0 or h2 <= 0:
            return 0

        cx1, cy1 = x11 + w1 / 2, y11 + h1 / 2
        cx2, cy2 = x21 + w2 / 2, y21 + h2 / 2
        dx = cx2 - cx1
        dy = cy2 - cy1
        dis = math.sqrt(dx * dx + dy * dy)
        if cnt <= 0 or self.videoFps <= 0:
            return 0

        dpix = 7.6 / ((w1 + w2) / 2) if (w1 + w2) > 0 else 0
        if dpix == 0:
            return 0

        v = dis * dpix * 3.6 / cnt * self.videoFps
        return round(v, 2)

    def calculateDirection(self, location1, location2):
        x11, y11, x12, y12 = location1
        x21, y21, x22, y22 = location2

        w1 = x12 - x11
        h1 = y12 - y11
        w2 = x22 - x21
        h2 = y22 - y21

        cx1, cy1 = x11 + w1 / 2, y11 + h1 / 2
        cx2, cy2 = x21 + w2 / 2, y21 + h2 / 2
        dx = cx2 - cx1
        dy = cy2 - cy1
        if dy > 0 and 3 * abs(dy) >= abs(dx):
            return 2
        if dy < 0 and 1.5 * abs(dy) >= abs(dx):
            return 0
        if dx > 0 and abs(dx) >= abs(dy):
            return 1
        if dx < 0 and abs(dx) >= abs(dy):
            return 3
        return 0

    def expand_bbox(self, x1, y1, x2, y2, frame_shape, ratio=0.3, padding=20):
        h, w = frame_shape[:2]
        expand_w = max(int((x2 - x1) * ratio), padding)
        expand_h = max(int((y2 - y1) * ratio), padding)

        x1_new = max(0, x1 - expand_w)
        x2_new = min(w, x2 + expand_w)
        y1_new = max(0, y1 - expand_h)
        y2_new = min(h, y2 + expand_h)

        return y1_new, y2_new, x1_new, x2_new

    def save_illegal_speed(self, id, x1, y1, x2, y2, frame):
        if id in self.uploaded_ids:
            return  # 已上传过该ID的违法信息，跳过
        if x1 >= x2 or y1 >= y2:
            print(f"[ERROR] 无效 BBox: x1={x1}, y1={y1}, x2={x2}, y2={y2}")
            return

        save_dir = os.path.join(
            "D:\GJ\Azuoye\Axsq\\third\A\\traffic_work\\vue-project1\public\images\\violations",
            "illegal", "overspeed"
        )
        os.makedirs(save_dir, exist_ok=True)

        try:
            # 使用扩展后的区域
            y1_new, y2_new, x1_new, x2_new = self.expand_bbox(x1, y1, x2, y2, frame.shape)

            img_tmp = frame[y1_new:y2_new, x1_new:x2_new]

            img_pil = Image.fromarray(cv2.cvtColor(img_tmp, cv2.COLOR_BGR2RGB))
            timestamp_str = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"{self.idx_frame}_{timestamp_str}_{self.carLicense.get(id, '')}.jpg"
            img_pil.save(os.path.join(save_dir, filename))

            # 新增：上传违法信息
            self.upload_illegal_info(
                illegal_type="overspeed",
                car_id=id,
                license_plate=self.carLicense.get(id, ""),
                image_path=os.path.join(save_dir, filename)
            )
            self.illegal[self.idx_frame] = {'overspeed': f'{id} {self.carLicense.get(id, "")}'}
            self.uploaded_ids.add(id)
        except Exception as e:
            print(f"[ERROR] 保存超速截图失败: {e}")

    def save_illegal_turnwrong(self, id, x1, y1, x2, y2, frame):
        if id in self.uploaded_ids:
            return  # 已上传过该ID的违法信息，跳过
        if x1 >= x2 or y1 >= y2:
            print(f"[ERROR] 无效 BBox: x1={x1}, y1={y1}, x2={x2}, y2={y2}")
            return

        save_dir = os.path.join(
            "D:\GJ\Azuoye\Axsq\\third\A\\traffic_work\\vue-project1\public\images\\violations",
            "illegal", "turnwrong"
        )
        os.makedirs(save_dir, exist_ok=True)

        try:
            y1_new, y2_new, x1_new, x2_new = self.expand_bbox(x1, y1, x2, y2, frame.shape)

            img_tmp = frame[y1_new:y2_new, x1_new:x2_new]

            img_pil = Image.fromarray(cv2.cvtColor(img_tmp, cv2.COLOR_BGR2RGB))
            timestamp_str = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"{self.idx_frame}_{timestamp_str}_{self.carLicense.get(id, '')}.jpg"
            img_pil.save(os.path.join(save_dir, filename))

            # 新增：上传违法信息
            self.upload_illegal_info(
                illegal_type="turnwrong",
                car_id=id,
                license_plate=self.carLicense.get(id, ""),
                image_path=os.path.join(save_dir, filename)
            )
            self.illegal[self.idx_frame] = {'turnwrong': f'{id} {self.carLicense.get(id, "")}'}
            self.uploaded_ids.add(id)
        except Exception as e:
            print(f"[ERROR] 保存未按导向行驶截图失败: {e}")

    def save_illegal_lane(self, id, x1, y1, x2, y2, frame):
        if id in self.uploaded_ids:
            return  # 已上传过该ID的违法信息，跳过
        save_dir = os.path.join(
            "D:\GJ\Azuoye\Axsq\\third\A\\traffic_work\\vue-project1\public\images\\violations",
            "illegal", "touchline"
        )
        os.makedirs(save_dir, exist_ok=True)

        try:
            y1_new, y2_new, x1_new, x2_new = self.expand_bbox(x1, y1, x2, y2, frame.shape)

            img_tmp = frame[y1_new:y2_new, x1_new:x2_new]

            img_pil = Image.fromarray(cv2.cvtColor(img_tmp, cv2.COLOR_BGR2RGB))
            timestamp_str = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"{self.idx_frame}_{timestamp_str}_{self.carLicense.get(id, '')}.jpg"
            img_pil.save(os.path.join(save_dir, filename))

            # 新增：上传违法信息
            self.upload_illegal_info(
                illegal_type="touchline",
                car_id=id,
                license_plate=self.carLicense.get(id, ""),
                image_path=os.path.join(save_dir, filename)
            )
            self.illegal[self.idx_frame] = {'touchline': f'{id} {self.carLicense.get(id, "")}'}
            self.uploaded_ids.add(id)
        except Exception as e:
            print(f"[ERROR] 保存压线截图失败: {e}")


    def carrunred(self, id, x1, x2, y1, y2, ori_im):
        if id in self.uploaded_ids:
            return  # 已上传过该ID的违法信息，跳过
        if (
                self.detector.trafficLightColor == 'red' or self.detector.forwardlightColor == 'red'
        ) and self.carDirection.get(id, "") == 'UP' and self.carSpeed.get(id, 0) >= 10:

            if not self.carFromForward.get(id, False):
                print("[DEBUG] 未进入直行区域，跳过闯红灯记录")
                return

            save_dir = os.path.join(
                "D:\GJ\Azuoye\Axsq\\third\A\\traffic_work\\vue-project1\public\images\\violations",
                "illegal", "carrunred"
            )
            os.makedirs(save_dir, exist_ok=True)

            try:
                y1_new, y2_new, x1_new, x2_new = self.expand_bbox(x1, y1, x2, y2, ori_im.shape)

                imgTmp = ori_im[y1_new:y2_new, x1_new:x2_new]

                timestamp_str = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"{self.idx_frame}_{timestamp_str}_{self.carLicense.get(id, '')}.jpg"
                cv2.imwrite(os.path.join(save_dir, filename))

                # 新增：上传违法信息
                self.upload_illegal_info(
                    illegal_type="carrunred",
                    car_id=id,
                    license_plate=self.carLicense.get(id, ""),
                    image_path=os.path.join(save_dir, filename)
                )
                self.illegal[self.idx_frame] = {'carrunred': f'{id} {self.carLicense.get(id, "")}'}
                print(f"[INFO] 成功保存闯红灯截图 ID={id}")
                self.uploaded_ids.add(id)
            except Exception as e:
                print(f"[ERROR] 保存闯红灯截图失败: {e}")

    def judge_line_illegal(self, id, car_bbox, config):
        if not config.laneline or len(config.laneline) < 4:
            return False

        arr = []
        for i in range(len(config.laneline)):
            arr.append(config.laneline[i])
            if (i + 1) % 4 == 0:
                if len(arr) < 4:
                    continue
                x1, y1, x2, y2 = arr
                arr = []
                q1 = [x1, y1]
                q2 = [x2, y2]
                tx1, ty1, tx2, ty2 = car_bbox
                p1 = [tx1, ty1]
                p2 = [tx2, ty2]
                if self.isIntersect(p1, p2, q1, q2):
                    return True
        return False

    def isIntersect(self, p1, p2, q1, q2):
        pass

    def inArea(self, x, y, area):
        x1, y1, x2, y2, x3, y3, x4, y4 = area
        if x > (x1 + x2) / 2 and x < (x3 + x4) / 2 and y > y1 and y < y2:
            return True
        return False


    def plot_lane(self, area, im, label=""):
        x1, y1, x2, y2, x3, y3, x4, y4 = area
        cv2.putText(im, label, (x1, y1 - 2), 0, 3, [225, 255, 255], thickness=2, lineType=cv2.LINE_AA)
        cv2.line(im, (x1, y1), (x2, y2), (255, 0, 0), 5)
        cv2.line(im, (x2, y2), (x4, y4), (255, 0, 0), 5)
        cv2.line(im, (x4, y4), (x3, y3), (255, 0, 0), 5)
        cv2.line(im, (x3, y3), (x1, y1), (255, 0, 0), 5)


    def emitImg(self, cvImg, sinImg):
        pass

    def endDetect(self):
        self.endFlag = True
        if self.vdo.isOpened():
            self.vdo.release()

    def upload_illegal_info(self, illegal_type, car_id, license_plate, image_path):
        """
        向 Spring Boot 后端上传违法信息
        :param illegal_type: str, 违法类型（如 "overspeed", "turnwrong", "runred" 等）
        :param car_id: int, 车辆 ID
        :param license_plate: str, 车牌号
        :param image_path: str, 截图本地路径
        """
        url = "http://localhost:8080/admin/violation/uploadIllegal"  # 替换为你的后端接口地址

        # 判断违法主体是人还是车
        subject = "person" if illegal_type == "runred" else "car"


        data = {
            "subject": subject,
            "illegalType": illegal_type,
            "licensePlate": license_plate,
            "screenshotPath": image_path
        }
        try:
            response = requests.post(url, json=data)
            if response.status_code == 200:
                print(f"[INFO] 成功上传违法信息 ID={car_id}, 类型={illegal_type}, 主体={subject}")
            else:
                print(f"[ERROR] 上传失败，状态码: {response.status_code}")
        except Exception as e:
            print(f"[ERROR] 上传违法信息失败: {e}")



